#!/usr/bin/env python3

class Node:
    def __init__(self, key, left=None, right=None, color=None, parent=None):
        self.key = key
        self.left = left
        self.right = right
        self.parent = parent
        self.color = color

    def is_red(self):
        return self.color == 'red'


class RedBlackTree:
    def __init__(self):
        self.root = None
        self.size = 0
        
    def insert(self, key):
        node = Node(key=key, color='red')
        if self.root is None:
            self.root = node
            self.root.color = 'black'
        else:
            self._bst_insert(node)
            self._fix_volitation(node) 

    def print_tree(self, showview=False):
        if self.root is None:
            print('RedBlack Tree is Empty')
            return
        from graphviz import Digraph
        dot = Digraph(comment='The Round Table')
        dot.node(str(self.root.key), color=self.root.color, shape='circle')
        self._print_tree_internal(self.root, dot)
        print(dot.source)
        if showview:
            dot.view()

    def _print_tree_internal(self, current_root, dot):
        if not current_root.left and not current_root.right:
            return
        left_name = None
        if current_root.left:
            left_name = str(current_root.left.key)
            dot.node(left_name, color=current_root.left.color, shape='circle')
            dot.edge(str(current_root.key), left_name)
            self._print_tree_internal(current_root.left, dot)
        else:
            left_name = str(hash(current_root)) + 'left'
            dot.node(left_name, shape='point')
            dot.edge(str(current_root.key), left_name)

        right_name = None
        if current_root.right:
            right_name = str(current_root.right.key)
            dot.node(right_name, color=current_root.right.color, shape='circle')
            dot.edge(str(current_root.key), right_name)
            self._print_tree_internal(current_root.right, dot)
        else:
            right_name = str(hash(current_root)) + 'right'
            dot.node(right_name, shape='point')
            dot.edge(str(current_root.key), right_name)
        

    def _fix_volitation(self, node):
        while node.parent and node.parent.color != 'black':
            print("node.key: ", node.key)
            self.print_tree()
            parent = node.parent
            grant_parent = parent.parent
            uncle = grant_parent.left if grant_parent.right == parent else grant_parent.right
            if uncle:
                print("uncle.key: ", uncle.key, ", and uncle.color: ", uncle.color)
            else:
                print("not found uncle")
            if uncle and uncle.color == 'red':
                parent.color = 'black'
                uncle.color = 'black'
                node = grant_parent
                node.color = 'red'
            else:
                if grant_parent.left == parent:
                    if node == parent.left:
                        # left left case
                        self._right_rotate(grant_parent, parent, True)
                    else:
                        # left right case
                        self._left_rotate(parent, node)
                        self._right_rotate(grant_parent, node, True)
                else:
                    if node == parent.right:
                        # right right case
                        if node.key == 5:
                            print("found left rotate")
                        self._left_rotate(grant_parent, parent, True)
                    else:
                        # right left case
                        self._right_rotate(parent, node)
                        self._left_rotate(grant_parent, node, True)

        if node == self.root and node.color == 'red':
            node.color = 'black'

    def _right_rotate(self, g, p, swapcolor=False):
        p.parent = g.parent
        g.left = p.right
        if p.right:
            p.right.parent = g
        p.right = g
        if g.parent:
            if g.parent.left == g:
                g.parent.left = p
            else:
                g.parent.right = p
        g.parent = p
        if g == self.root:
            self.root = p
        if swapcolor:
            color = p.color
            p.color = g.color
            g.color = color

    def _left_rotate(self, g, p, swapcolor=False):
        p.parent = g.parent
        g.right = p.left
        if p.left:
            p.left.parent = g
        p.left = g
        if g.parent:
            if g.parent.left == g:
                g.parent.left = p
            else:
                g.parent.right = p
        g.parent = p
        if g == self.root:
            self.root = p
        if swapcolor:
            color = p.color
            p.color = g.color
            g.color = color

    def _bst_insert(self, node):
        current = self.root
        while True:
            if current.key < node.key:
                if current.right:
                    current = current.right
                else:
                    node.parent = current
                    current.right = node
                    break
            else:
                if current.left:
                    current = current.left
                else:
                    node.parent = current
                    current.left = node
                    break
                    


def insert():
    rbtree = RedBlackTree()
#    for index in range(1, 100): 
#        rbtree.insert(index)
    for key in [3, 7, 18, 10, 22, 8, 11, 26]:
        rbtree.insert(key)
    rbtree.print_tree(True)

insert()
